# grid_search.py
import json
import os
import subprocess
import itertools
import pandas as pd
from datetime import datetime
import argparse
import shutil

def run_training_subprocess(base_command, params_dict, run_output_dir):
    """构建并运行 train_model.py 的子进程。"""
    command = list(base_command) # 创建副本以避免修改原始列表
    
    # 添加 output_dir 参数，指向特定于此参数组合的目录
    command.extend(['--output_dir', run_output_dir])

    for key, value in params_dict.items():
        if value is None and key == 'rf_max_depth': # 特殊处理 rf_max_depth=None
             # train_model.py 的 argparse 会将 None 处理为字符串 "None"
             # 对于 rf_max_depth，我们希望不传递该参数，或者传递一个能被 int(None) 接受的值
             # 或者修改 train_model.py 来正确处理字符串 "None"
             # 这里我们选择不传递它，让 argparse 使用其默认值 (通常是 None)
             # 但如果 train_model.py 的默认值不是 None，这可能会有问题。
             # 最好的方式是在 train_model.py 中处理 "None" 字符串。
             # 假设 train_model.py 的 argparse 对 --rf_max_depth 的 type=int，default=None
             # 传递 "None" 字符串会导致 argparse 错误。
             # 如果参数的默认值就是我们想要的 None，我们可以跳过添加它。
             # 然而，如果我们要明确设置它为 None，而 argparse 的默认值不同，这就复杂了。
             # 一个简单的解决方法是，如果值为 None 且是特定参数，就不添加到命令中。
             # 这依赖于 train_model.py 中 argparse 的 default 设置。
             # 为了更明确，我们将传递一个特殊的值，让 train_model.py 去解析。
             # 或者，如果 train_model.py 的 default=None，我们就不添加这个参数。
             # 假设 train_model.py 的 default for rf_max_depth is None.
             # 如果 value is None，我们就不添加 --rf_max_depth 参数。
             pass # Skip adding --rf_max_depth if value is None
        elif isinstance(value, bool):
            if value: # For boolean flags like --verbose
                command.append(f'--{key}')
        else:
            command.extend([f'--{key}', str(value)])
    
    print(f"  执行命令: {' '.join(command)}")
    
    try:
        # 将子进程的输出重定向到文件，以便调试
        stdout_log_path = os.path.join(run_output_dir, "stdout.log")
        stderr_log_path = os.path.join(run_output_dir, "stderr.log")

        with open(stdout_log_path, 'w') as f_stdout, open(stderr_log_path, 'w') as f_stderr:
            process = subprocess.Popen(command, stdout=f_stdout, stderr=f_stderr)
            process.wait() # 等待子进程完成

        if process.returncode == 0:
            print(f"  运行成功完成。日志: {stdout_log_path}, {stderr_log_path}")
            # 构造预期的平均指标文件名
            model_type = params_dict.get('model_type', 'unknown_model')
            avg_metrics_filename = f"{model_type}_avg_cv_metrics.csv"
            avg_metrics_path = os.path.join(run_output_dir, avg_metrics_filename)
            if os.path.exists(avg_metrics_path):
                return avg_metrics_path
            else:
                print(f"  错误: 平均指标文件未找到于 {avg_metrics_path}")
                return None
        else:
            print(f"  运行失败，返回码: {process.returncode}。查看日志: {stderr_log_path}")
            return None
    except Exception as e:
        print(f"  执行命令时发生异常: {e}")
        return None

def main(params_json_path):
    """主函数，用于执行网格搜索。"""
    try:
        with open(params_json_path, 'r') as f:
            params_config = json.load(f)
    except FileNotFoundError:
        print(f"错误: 参数文件 '{params_json_path}' 未找到。")
        return
    except json.JSONDecodeError:
        print(f"错误: 解析参数文件 '{params_json_path}' 失败。请检查JSON格式。")
        return

    common_params = params_config.get("common_params", {})
    models_to_search = params_config.get("models_to_search", {})
    grid_search_output_base = params_config.get("grid_search_output_base", "grid_search_results")

    if not os.path.exists(grid_search_output_base):
        os.makedirs(grid_search_output_base)

    all_results = []
    
    # train_model.py 脚本的路径 (假设在同一目录)
    train_script_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "train_model.py")
    if not os.path.exists(train_script_path):
        print(f"错误: train_model.py 未在预期的位置找到: {train_script_path}")
        print("请确保 train_model.py 与 grid_search.py 在同一目录下。")
        return

    base_command = ['python', train_script_path]

    for model_type, hyperparams_grid in models_to_search.items():
        print(f"\n--- 开始网格搜索: {model_type.upper()} ---")
        
        # 为当前模型类型创建一个子目录
        model_type_output_dir = os.path.join(grid_search_output_base, model_type)
        if not os.path.exists(model_type_output_dir):
            os.makedirs(model_type_output_dir)

        # 获取超参数的名称和它们的值的列表
        param_names = list(hyperparams_grid.keys())
        param_values_lists = [hyperparams_grid[name] for name in param_names]

        # 生成所有参数组合
        for i, param_combination_values in enumerate(itertools.product(*param_values_lists)):
            current_run_params = common_params.copy()
            current_run_params['model_type'] = model_type
            
            param_combination_dict = {}
            for name, value in zip(param_names, param_combination_values):
                current_run_params[name] = value
                param_combination_dict[name] = value
            
            # 为每个参数组合创建一个唯一的目录名
            # 使用参数值创建目录名，确保文件名安全
            param_str_parts = []
            for name, value in sorted(param_combination_dict.items()): # Sort for consistency
                # 移除参数名前缀 (如 svm_, rf_) 以简化目录名
                short_name = name.split(f"{model_type}_", 1)[-1] if name.startswith(f"{model_type}_") else name
                param_str_parts.append(f"{short_name}_{str(value).replace('.', 'p')}") # Replace . with p for float
            
            run_id = f"run_{i+1}_{'_'.join(param_str_parts)}"
            # 限制目录名长度 (如果需要)
            max_len = 100 # 示例最大长度
            if len(run_id) > max_len:
                run_id = run_id[:max_len] + "_etc"
            
            run_output_dir = os.path.join(model_type_output_dir, run_id)
            if not os.path.exists(run_output_dir):
                os.makedirs(run_output_dir)
            else:
                # 如果目录已存在，可以选择跳过或覆盖。这里我们选择覆盖（通过清空）。
                print(f"警告: 目录 {run_output_dir} 已存在。将清空并重新运行。")
                shutil.rmtree(run_output_dir)
                os.makedirs(run_output_dir)


            print(f"\n运行 {model_type.upper()} 第 {i+1} 次参数组合:")
            print(f"  参数: {param_combination_dict}")
            print(f"  输出目录: {run_output_dir}")

            avg_metrics_csv_path = run_training_subprocess(base_command, current_run_params, run_output_dir)

            if avg_metrics_csv_path and os.path.exists(avg_metrics_csv_path):
                try:
                    metrics_df = pd.read_csv(avg_metrics_csv_path)
                    if not metrics_df.empty:
                        # 将参数组合和指标合并
                        result_row = param_combination_dict.copy()
                        result_row['model_type'] = model_type
                        result_row['run_output_dir'] = run_output_dir 
                        for col in metrics_df.columns:
                            result_row[col] = metrics_df.iloc[0][col]
                        
                        # 检查是否有阈值搜索结果
                        threshold_path = os.path.join(run_output_dir, f"{model_type}_threshold_search.csv")
                        if os.path.exists(threshold_path):
                            try:
                                threshold_df = pd.read_csv(threshold_path)
                                # 找出F1-score最高的阈值
                                best_threshold_row = threshold_df.loc[threshold_df['f1_score'].idxmax()]
                                result_row['Best_Threshold'] = best_threshold_row['threshold']
                                result_row['Best_F1_Score'] = best_threshold_row['f1_score']
                                print(f"  最佳阈值: {best_threshold_row['threshold']:.2f}, F1-score: {best_threshold_row['f1_score']:.4f}")
                            except Exception as e:
                                print(f"  读取阈值搜索文件 {threshold_path} 时出错: {e}")
                        
                        all_results.append(result_row)
                        print(f"  成功获取指标: {metrics_df.iloc[0].to_dict()}")
                    else:
                        print(f"  警告: 指标文件 {avg_metrics_csv_path} 为空。")
                except Exception as e:
                    print(f"  读取指标文件 {avg_metrics_csv_path} 时出错: {e}")
            else:
                print(f"  未能获取 {model_type.upper()} 第 {i+1} 次运行的指标。")
                # 记录失败的运行
                failed_run_info = param_combination_dict.copy()
                failed_run_info['model_type'] = model_type
                failed_run_info['status'] = 'Failed'
                failed_run_info['run_output_dir'] = run_output_dir
                all_results.append(failed_run_info)


    # 将所有结果保存到一个汇总的CSV文件中
    if all_results:
        summary_df = pd.DataFrame(all_results)
        # 根据 Best_F1_Score (如果存在) 或 ROC-AUC 或 Accuracy 降序排列
        sort_key = 'Best_F1_Score' if 'Best_F1_Score' in summary_df.columns else 'ROC-AUC' if 'ROC-AUC' in summary_df.columns else 'Accuracy'
        if sort_key in summary_df.columns:
            summary_df = summary_df.sort_values(by=sort_key, ascending=False, na_position='last')
        
        summary_filename = f"grid_search_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
        summary_path = os.path.join(grid_search_output_base, summary_filename)
        summary_df.to_csv(summary_path, index=False)
        print(f"\n--- 网格搜索完成 ---")
        print(f"所有参数组合的性能指标已汇总到: {summary_path}")
        print("\n最佳参数组合 (基于主要指标排序):")
        print(summary_df.head())
    else:
        print("\n--- 网格搜索完成 ---")
        print("没有成功运行或没有结果可汇总。")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="执行基于JSON配置的网格搜索脚本。")
    parser.add_argument("--params_json", type=str, default="grid_search_params.json",
                        help="包含网格搜索参数的JSON配置文件路径。")
    args = parser.parse_args()
    
    main(args.params_json)
